package com.jobcn.spider;

import org.apache.commons.io.IOUtils;
import org.apache.http.ConnectionClosedException;
import org.apache.http.HttpResponse;
import org.apache.http.annotation.Contract;
import org.apache.http.annotation.ThreadingBehavior;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import us.codecraft.webmagic.Page;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Site;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.downloader.AbstractDownloader;
import us.codecraft.webmagic.downloader.HttpClientGenerator;
import us.codecraft.webmagic.downloader.HttpClientRequestContext;
import us.codecraft.webmagic.downloader.HttpUriRequestConverter;
import us.codecraft.webmagic.proxy.Proxy;
import us.codecraft.webmagic.proxy.ProxyProvider;
import us.codecraft.webmagic.selector.PlainText;
import us.codecraft.webmagic.utils.CharsetUtils;
import us.codecraft.webmagic.utils.HttpClientUtils;

import javax.net.ssl.SSLException;
import javax.net.ssl.SSLHandshakeException;
import java.io.IOException;
import java.net.SocketTimeoutException;
import java.nio.charset.Charset;
import java.util.HashMap;
import java.util.Map;

/**
 * 这是页面的下载器
 */
@Contract(threading = ThreadingBehavior.SAFE)
public class Downloader extends AbstractDownloader {

    private Logger logger = LoggerFactory.getLogger(getClass());

    //定制化------------start

    //统计成功次数，定时重启Httpclient
    //private static AtomicInteger downloadTimes = new AtomicInteger(0);
    //成功次数上限设为100
    //private static final int limitTimes = 1000;

    private static Map<Integer, String> errorCode = new HashMap<Integer, String>();

    static {
        errorCode.put(404, "Page not found! status code: 404 ");
        errorCode.put(407, "Get page error, try later! status code: 404 ");
        errorCode.put(429, "Too many request,try later! status code: 429");
        errorCode.put(503, "Proxy time out,try later! status code: 503");
    }
    //end--------------------------

    private final Map<String, CloseableHttpClient> httpClients = new HashMap<String, CloseableHttpClient>();

    private HttpClientGenerator httpClientGenerator = new HttpClientGenerator();

    private HttpUriRequestConverter httpUriRequestConverter = new HttpUriRequestConverter();
    
    private ProxyProvider proxyProvider;

    private boolean responseHeader = true;

    public void setHttpUriRequestConverter(HttpUriRequestConverter httpUriRequestConverter) {
        this.httpUriRequestConverter = httpUriRequestConverter;
    }

    public void setProxyProvider(ProxyProvider proxyProvider) {
        this.proxyProvider = proxyProvider;
    }

    private CloseableHttpClient getHttpClient(Site site) {
        if (site == null) {
            return httpClientGenerator.getClient(null);
        }
        String domain = site.getDomain();
        CloseableHttpClient httpClient = httpClients.get(domain);

        synchronized (this) {
            // 修改这里，重新创建httpclient,有些网站识别cookies,如果不禁用cookies可能需要定期刷新httpClient
            //if (downloadTimes.get() > limitTimes) {
            //    logger.warn("download {} tims,restart httpclient", downloadTimes);
            //    httpClient = null;
            //    httpClients.put(domain, httpClient);
            //    downloadTimes.set(0);
            //}
            if (httpClient == null) {
                httpClient = httpClients.get(domain);
                if (httpClient == null) {
                    httpClient = httpClientGenerator.getClient(site);
                    httpClients.put(domain, httpClient);
                }
            }
        }
        return httpClient;
    }

    @Override
    protected void onSuccess(Request request) {
        //成功下载一次计数，每个httpclient实例限制爬取页数
        //downloadTimes.getAndIncrement();
    }

    @Override
    public Page download(Request request, Task task) {
        if (task == null || task.getSite() == null) {
            throw new NullPointerException("task or site can not be null");
        }
//        logger.debug("downloading page {}", request.getUrl());
        CloseableHttpResponse httpResponse = null;
        CloseableHttpClient httpClient = getHttpClient(task.getSite());
        Proxy proxy = proxyProvider != null ? proxyProvider.getProxy(task) : null;
        HttpClientRequestContext requestContext = httpUriRequestConverter.convert(request, task.getSite(), proxy);
        Page page = Page.fail();
        try {
            httpResponse = httpClient.execute(requestContext.getHttpUriRequest(), requestContext.getHttpClientContext());
            page = handleResponse(request, task.getSite().getCharset(), httpResponse, task);
            onSuccess(request);
            logger.debug("downloading page success {}", request.getUrl());
//            return page;
        } catch (ConnectionClosedException e) {
            logger.warn("Premature end of chunk coded message body: {}", request.getUrl());
            onError(request);
        } catch (SSLHandshakeException e) {
            logger.warn("Remote host closed connection during handshake: {}", request.getUrl());
            onError(request);
        } catch (SSLException e) {
            logger.warn("SSL peer shut down incorrectly: {}", request.getUrl());
            onError(request);
        } catch (SocketTimeoutException e) {
            logger.warn("download page time out:{}", request.getUrl());
            onError(request);
        } catch (IOException e) {
            logger.warn("download page error:{} ", request.getUrl());
            onError(request);
        }finally {
            if (httpResponse != null) {
                //ensure the connection is released back to pool
                EntityUtils.consumeQuietly(httpResponse.getEntity());
            }
            if (proxyProvider != null && proxy != null) {
                proxyProvider.returnProxy(proxy, page, task);
            }
        }
        //如果下载成功了，并不代表网页正常，这里要再做处理
        if (page.isDownloadSuccess()) {
            //获得请求返回的状态码
            int statusCode = httpResponse.getStatusLine().getStatusCode();
            if (errorCode.containsKey(statusCode)) {
                logger.warn(errorCode.get(statusCode));
                page.setDownloadSuccess(false);
                //page.getRequest().putExtra(Request., false);
            }
        }
        return page;
    }

    @Override
    public void setThread(int thread) {
        httpClientGenerator.setPoolSize(thread);
    }

    protected Page handleResponse(Request request, String charset, HttpResponse httpResponse, Task task) throws IOException {
        String content = getResponseContent(charset, httpResponse);
        Page page = new Page();
        page.setRawText(content);
        page.setUrl(new PlainText(request.getUrl()));
        page.setRequest(request);
        page.setStatusCode(httpResponse.getStatusLine().getStatusCode());
        page.setDownloadSuccess(true);
        if (responseHeader) {
            page.setHeaders(HttpClientUtils.convertHeaders(httpResponse.getAllHeaders()));
        }
        return page;
    }

    private String getResponseContent(String charset, HttpResponse httpResponse) throws IOException {
        if (charset == null) {
            byte[] contentBytes = IOUtils.toByteArray(httpResponse.getEntity().getContent());
            String htmlCharset = getHtmlCharset(httpResponse, contentBytes);
            if (htmlCharset != null) {
                return new String(contentBytes, htmlCharset);
            } else {
                logger.warn("Charset autodetect failed, use {} as charset. Please specify charset in Site.setCharset()", Charset.defaultCharset());
                return new String(contentBytes);
            }
        } else {
            return IOUtils.toString(httpResponse.getEntity().getContent(), charset);
        }
    }

    private String getHtmlCharset(HttpResponse httpResponse, byte[] contentBytes) throws IOException {
        return CharsetUtils.detectCharset(httpResponse.getEntity().getContentType() == null ? "" : httpResponse.getEntity().getContentType().getValue(), contentBytes);
    }
}
